from typing import Any, Dict, List, Optional, Tuple

from chernc.agents.apipool import AsyncLLMAPIPool, InputLengthScorer
from chernc.llm.lm_inference import AsyncLanguageModelPredictor, LanguageModelPredictor

rules_for_kind = {
    # "GlobalVariable": [
    #     "global variables are preferred to be `static mut`.",
    # ]
}

TRANSLATION_MAX_TOKENS = 8192  # 模型输出的最大token数量限制


class TranslatorBase(object):
    def __init__(self) -> None:
        pass

    async def translate(
        self,
        source: str,
        deps: str,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        raise NotImplementedError()

    async def fix(
        self,
        to_fix: str,
        deps: str,
        err_msgs: str,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        raise NotImplementedError()


class Translator(TranslatorBase):
    """
    A class to translate code from one language to another.
    """

    def __init__(self, inference_engine: LanguageModelPredictor) -> None:
        self.inference_engine: LanguageModelPredictor = inference_engine

    def translate(
        self,
        source_code: str,
        source_lang: str = "cpp",
        target_lang: str = "rust",
        few_shot_examples: Optional[List[Dict]] = None,
        max_tokens: int = 4096,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": f"Please translate the following {source_lang} code into {target_lang}:"},
            {"role": "assistant", "content": "OK, please provide the source code, I will translate the code."},
        ]
        if few_shot_examples is not None:
            messages.extend(few_shot_examples)
        messages.append({"role": "user", "content": source_code})
        return self.inference_engine.chat(messages=messages, max_length=max_tokens, temperature=temperature)

    def translate_chunk_to_rust(
        self,
        chunk_source: str,
        chunk_deps: Optional[List[Tuple[str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, dtarget in chunk_deps:
                messages.append({"role": "user", "content": dsource})
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append({"role": "user", "content": chunk_source})

        llm_res = self.inference_engine.chat(messages=messages, max_length=max_tokens, temperature=temperature)
        return llm_res

    def make_source_promt(self, source, decls=None):
        if not decls and len(decls) != 0:
            return {
                "role": "user",
                "content": f"The code is :\n {source} \n The summury of declaration(s) in this code snippet is: \n {decls}",
            }
        else:
            return {"role": "user", "content": source}

    def translate_chunk_and_decls_to_rust(
        self,
        chunk_source: str,
        chunk_decls: str,
        chunk_deps: Optional[List[Tuple[str, str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, decls, dtarget in chunk_deps:
                messages.append(self.make_source_promt(dsource, decls))
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append(self.make_source_promt(dsource, chunk_decls))

        llm_res = self.inference_engine.chat(messages=messages, max_length=max_tokens, temperature=temperature)
        return llm_res


class AsyncTranslator(TranslatorBase):
    """
    A class to translate code from one language to another.
    """

    def __init__(self, inference_engine: AsyncLanguageModelPredictor) -> None:
        self.inference_engine: AsyncLanguageModelPredictor = inference_engine

    async def translate_chunk_to_rust(
        self,
        chunk_source: str,
        chunk_deps: Optional[List[Tuple[str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent.  I'll provide some declarations information of the snippet. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, dtarget in chunk_deps:
                messages.append({"role": "user", "content": dsource})
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append({"role": "user", "content": chunk_source})

        llm_res = await self.inference_engine.chat(messages=messages, max_length=max_tokens, temperature=temperature)
        return llm_res

    def make_source_promt(self, source, decls=None):
        if not decls and len(decls) != 0:
            return {
                "role": "user",
                "content": f"The code is :\n {source} \n The summury of declaration(s) in this code snippet is: \n {decls}",
            }
        else:
            return {"role": "user", "content": source}

    async def translate_chunk_and_decls_to_rust(
        self,
        chunk_source: str,
        chunk_decls: str,
        chunk_deps: Optional[List[Tuple[str, str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent. I'll provide some declarations information of the snippet. Note that C global variable should be translated to Rust static variable, but global functions do not need a static declaration, they are just oridinary fn. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, decls, dtarget in chunk_deps:
                messages.append(self.make_source_promt(dsource, decls))
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append(self.make_source_promt(chunk_source, chunk_decls))

        llm_res = await self.inference_engine.chat(messages=messages, max_length=max_tokens, temperature=temperature)
        return llm_res


class PoolBasedModelTranslator(TranslatorBase):
    def __init__(self) -> None:
        self.api_pool = AsyncLLMAPIPool()

    def register_a_model(self, name: str, predicator: AsyncLanguageModelPredictor, model_features: Dict[str, Any]):
        self.api_pool.register_model(name, predicator, model_features)

    def calculate_message_length(self, messages):
        total_length = sum(len(key) + len(value) for d in messages for key, value in d.items())
        return total_length

    async def translate_chunk_to_rust(
        self,
        chunk_source: str,
        chunk_deps: Optional[List[Tuple[str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent.  I'll provide some declarations information of the snippet. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, dtarget in chunk_deps:
                messages.append({"role": "user", "content": dsource})
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append({"role": "user", "content": chunk_source})

        llm_res = await self.api_pool.call_chat(
            messages=messages,
            selectors=[InputLengthScorer()],
            max_length=max_tokens,
            temperature=temperature,
        )
        return llm_res

    def make_source_promt(self, source, decls=None):
        if not decls and len(decls) != 0:
            return {
                "role": "user",
                "content": f"The code is :\n {source} \n The summury of declaration(s) in this code snippet is: \n {decls}",
            }
        else:
            return {"role": "user", "content": source}

    async def translate_chunk_and_decls_to_rust(
        self,
        chunk_source: str,
        chunk_decls: str,
        chunk_deps: Optional[List[Tuple[str, str, str]]] = None,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        messages = [
            {"role": "system", "content": "You are a helpful assistant."},
            {
                "role": "user",
                "content": f"Please translate the following c/cpp code snippets into rust code. these snippets comes from one project, so I want you to make sure that translation results are consistent. I'll provide some declarations information of the snippet. Note that C global variable should be translated to Rust static variable, but global functions do not need a static declaration, they are just oridinary fn. Moreover, don't complement unnecessary information, only translate provided parts",
            },
            {
                "role": "assistant",
                "content": "OK, please provide source code snippets, I will translate the code and make sure the translation of these code snippets are consistent.",
            },
        ]
        if chunk_deps:
            for dsource, decls, dtarget in chunk_deps:
                messages.append(self.make_source_promt(dsource, decls))
                messages.append({"role": "assistant", "content": f"Ok, the translation result is ```rust\n {dtarget}\n``` "})
        messages.append(self.make_source_promt(chunk_source, chunk_decls))

        llm_res = await self.api_pool.call_chat(
            {"exp_len": self.calculate_message_length(messages)},
            messages=messages,
            max_length=max_tokens,
            temperature=temperature,
        )
        return llm_res

    async def translate(
        self,
        source: str,
        deps: str,
        # kinds: List[str],
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.1,
    ):
        rules = [
            # "keep every CHERNC_MARK and CHERNC_END_MARK to surround the transpilation result, just like:\n```rust\n// CHERNC_MARK: /path/to/file.c:1\nfn foo(){\n    // some code}\n// CHERNC_END_MARK: /path/to/file.c:2\n```",
        ]
        user_content = "Please translate a C/C++ top-level snippet into Rust. \n"
        if rules:
            user_content += "Notice that:\n"
            for rule in rules:
                user_content += f"- {rule}\n"
        # for decl in kinds:
        #     for rule in rules_for_kind.get(decl, []):
        #         user_content += f"- {rule}\n"

        user_content += "The C/C++ snippet to transpile:\n```\n" + source + "\n```\n"
        if deps:
            user_content += "The Rust context:\n```rust\n"
            user_content += deps
            user_content += "```\n"

        messages = [
            {
                "role": "system",
                "content": "You are an expert programming assistant familiar with C/C++ and Rust, helping to transcompile a C/C++ snippet into Rust.",
            },
            {"role": "user", "content": user_content},
        ]

        llm_res = await self.api_pool.call_chat(
            messages=messages,
            selectors=[InputLengthScorer()],
            max_length=max_tokens,
            temperature=temperature,
        )
        return llm_res

    async def fix(
        self,
        to_fix: str,
        deps: str,
        err_msgs: str,
        max_tokens: int = TRANSLATION_MAX_TOKENS,
        temperature: float = 0.0,
    ):
        rules = [
            # "keep every CHERNC_MARK and CHERNC_END_MARK aside its surrounded code.",
            "**DO NOT** try to revise the code outside this buggy snippet. ",
        ]
        user_content = "Please fix a Rust snippet according to the compilation error messages. This snippet is sliced from a Rust file. And notice that:\n"
        for rule in rules:
            user_content += f"- {rule}\n"
        # for decl in kinds:
        #     for rule in rules_for_kind.get(decl, []):
        #         user_content += f"- {rule}\n"

        user_content += "The buggy Rust snippet to fix:\n```rust\n" + to_fix + "\n```\n"
        user_content += (
            "Compilation errors related to the buggy snippet, "
            "where line numbers is calculated relative to this snippet, and zero or negative means outside this snippet:\n"
        )
        user_content += err_msgs + "\n"
        if deps:
            user_content += "The Rust context before the buggy snippet:\n```rust\n"
            user_content += deps
            user_content += "```\n"
        messages = [
            {
                "role": "system",
                "content": "You are an expert programming assistant familiar with Rust, helping to fix a buggy Rust snippet.",
            },
            {"role": "user", "content": user_content},
        ]

        llm_res = await self.api_pool.call_chat(
            messages=messages,
            selectors=[InputLengthScorer()],
            max_length=max_tokens,
            temperature=temperature,
        )
        return llm_res
